/*
 * (c) Copyright 2021 Palantir Technologies Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.palantir.atlasdb.timelock;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.MoreCollectors;
import com.palantir.atlasdb.config.AtlasDbRuntimeConfig;
import com.palantir.atlasdb.config.ImmutableAtlasDbConfig;
import com.palantir.atlasdb.keyvalue.api.LockWatchCachingConfig;
import com.palantir.atlasdb.keyvalue.api.Namespace;
import com.palantir.atlasdb.keyvalue.api.TableReference;
import com.palantir.atlasdb.keyvalue.api.watch.LockWatchManagerInternal;
import com.palantir.atlasdb.table.description.Schema;
import com.palantir.atlasdb.table.description.TableDefinition;
import com.palantir.atlasdb.table.description.ValueType;
import com.palantir.atlasdb.transaction.api.ConflictHandler;
import com.palantir.atlasdb.transaction.api.PreCommitCondition;
import com.palantir.atlasdb.transaction.api.TransactionManager;
import com.palantir.lock.watch.LockEvent;
import com.palantir.lock.watch.LockRequestMetadata;
import com.palantir.lock.watch.LockWatchCreatedEvent;
import com.palantir.lock.watch.LockWatchEvent;
import com.palantir.lock.watch.LockWatchReferences;
import com.palantir.lock.watch.LockWatchReferences.LockWatchReference;
import com.palantir.lock.watch.UnlockEvent;
import com.palantir.logsafe.Preconditions;
import java.time.Duration;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import org.awaitility.Awaitility;

public final class LockWatchIntegrationTestUtilities {
    private static final String TEST_PACKAGE = "package";
    static final String TABLE = "table";
    static final String TABLE_2 = "rowLevelTable";

    private LockWatchIntegrationTestUtilities() {
        // no-op
    }

    /**
     * By requesting a snapshot, we get a single snapshot event that contains all currently held lock descriptors.
     * As unlocks are asynchronous, we must wait for the unlocks to go through to guarantee consistent testing.
     */
    public static void awaitAllUnlocked(TransactionManager txnManager) {
        LockWatchManagerInternal lockWatchManager = extractInternalLockWatchManager(txnManager);
        Awaitility.await("All descriptors are unlocked")
                .atMost(Duration.ofSeconds(5))
                .pollDelay(Duration.ofMillis(100))
                .until(() -> getLockWatchState(txnManager, lockWatchManager)
                        .map(event -> event.accept(new LockWatchCreatedEventVisitor(
                                createdEvent -> createdEvent.lockDescriptors().isEmpty())))
                        .orElse(false));
    }

    /**
     * The lock watch manager registers watch events every five seconds - therefore, tables may not be watched
     * immediately after a TimeLock leader election. By requesting an update with an empty version, we guarantee that
     * we only receive a single event containing a snapshot of the current lock watch event state, including watched
     * tables.
     */
    public static void awaitTableWatched(TransactionManager txnManager, TableReference tableReference) {
        LockWatchManagerInternal lockWatchManager = extractInternalLockWatchManager(txnManager);
        Awaitility.await("Tables are watched")
                .atMost(Duration.ofSeconds(5))
                .pollDelay(Duration.ofMillis(100))
                .until(() -> getLockWatchState(txnManager, lockWatchManager)
                        .map(event -> event.accept(new LockWatchCreatedEventVisitor(createdEvent -> createdEvent
                                .references()
                                .contains(LockWatchReferences.entireTable(tableReference.getQualifiedName())))))
                        .orElse(false));
    }

    public static void awaitLockWatchCreated(TransactionManager txnManager, LockWatchReference lockWatchReference) {
        LockWatchManagerInternal lockWatchManager = extractInternalLockWatchManager(txnManager);
        Awaitility.await("Lock watch created event is fired")
                .atMost(Duration.ofSeconds(5))
                .pollDelay(Duration.ofMillis(100))
                .until(() -> getLockWatchEvents(txnManager, lockWatchManager).stream()
                        .anyMatch(event -> event.accept(new LockWatchCreatedEventVisitor(
                                createdEvent -> createdEvent.references().contains(lockWatchReference)))));
    }

    /**
     * The internal version of the lock watch manager is hidden from the user, both to reduce API surface area, and
     * because certain classes aren't visible everywhere.
     */
    public static LockWatchManagerInternal extractInternalLockWatchManager(TransactionManager txnManager) {
        return (LockWatchManagerInternal) txnManager.getLockWatchManager();
    }

    public static List<Optional<LockRequestMetadata>> extractMetadata(List<LockWatchEvent> lockWatchEvents) {
        return lockWatchEvents.stream()
                .map(event -> event.accept(LockEventVisitor.INSTANCE))
                .filter(Optional::isPresent)
                .map(Optional::get)
                .map(LockEvent::metadata)
                .collect(Collectors.toList());
    }

    public static TransactionManager createTransactionManager(
            double validationProbability, TestableTimelockCluster timelockCluster, String namespace) {
        return TimeLockTestUtils.createTransactionManager(
                        timelockCluster,
                        namespace,
                        AtlasDbRuntimeConfig.defaultRuntimeConfig(),
                        ImmutableAtlasDbConfig.builder()
                                .lockWatchCaching(LockWatchCachingConfig.builder()
                                        .validationProbability(validationProbability)
                                        .build()),
                        Optional.empty(),
                        createSchema())
                .transactionManager();
    }

    private static Optional<LockWatchEvent> getLockWatchState(
            TransactionManager txnManager, LockWatchManagerInternal lockWatchManager) {
        return txnManager.runTaskThrowOnConflict(txn -> lockWatchManager
                .getCache()
                .getEventCache()
                .getUpdateForTransactions(ImmutableSet.of(txn.getTimestamp()), Optional.empty())
                .events()
                .stream()
                .collect(MoreCollectors.toOptional()));
    }

    private static List<LockWatchEvent> getLockWatchEvents(
            TransactionManager txnManager, LockWatchManagerInternal lockWatchManager) {
        return txnManager.runTaskThrowOnConflict(txn -> lockWatchManager
                .getCache()
                .getEventCache()
                .getUpdateForTransactions(ImmutableSet.of(txn.getTimestamp()), Optional.empty())
                .events());
    }

    private static Schema createSchema() {
        Schema schema = new Schema("table", TEST_PACKAGE, Namespace.DEFAULT_NAMESPACE);

        TableDefinition tableDef = new TableDefinition() {
            {
                rowName();
                rowComponent("key", ValueType.BLOB);
                noColumns();
                enableCaching();
                conflictHandler(ConflictHandler.SERIALIZABLE_CELL);
            }
        };
        schema.addTableDefinition(TABLE, tableDef);

        TableDefinition tableWithoutCaching = new TableDefinition() {
            {
                rowName();
                rowComponent("key", ValueType.BLOB);
                noColumns();
                conflictHandler(ConflictHandler.SERIALIZABLE_CELL);
            }
        };
        schema.addTableDefinition(TABLE_2, tableWithoutCaching);

        return schema;
    }

    private static final class LockWatchCreatedEventVisitor implements LockWatchEvent.Visitor<Boolean> {
        private final Predicate<LockWatchCreatedEvent> predicate;

        private LockWatchCreatedEventVisitor(Predicate<LockWatchCreatedEvent> predicate) {
            this.predicate = predicate;
        }

        @Override
        public Boolean visit(LockEvent lockEvent) {
            return false;
        }

        @Override
        public Boolean visit(UnlockEvent unlockEvent) {
            return false;
        }

        @Override
        public Boolean visit(LockWatchCreatedEvent lockWatchCreatedEvent) {
            return predicate.test(lockWatchCreatedEvent);
        }
    }

    private enum LockEventVisitor implements LockWatchEvent.Visitor<Optional<LockEvent>> {
        INSTANCE;

        @Override
        public Optional<LockEvent> visit(LockEvent lockEvent) {
            return Optional.of(lockEvent);
        }

        @Override
        public Optional<LockEvent> visit(UnlockEvent unlockEvent) {
            return Optional.empty();
        }

        @Override
        public Optional<LockEvent> visit(LockWatchCreatedEvent lockWatchCreatedEvent) {
            return Optional.empty();
        }
    }

    /**
     * {@link PreCommitCondition} is actually run several times throughout a transaction - namely, before any read
     * if this is configured. When run at these times, the start timestamp is passed in. However, when the condition
     * is run at commit time, the commit timestamp is actually passed in, which means that we have a way of determining
     * whether the condition is being evaluated at commit time or not.
     *
     * Given that this condition may execute arbitrary code, this is one way to guarantee that test code will be
     * run during the commit flow. To achieve this, we set the start timestamp, then evaluate each invocation of
     * {@link PreCommitCondition#throwIfConditionInvalid(long)} against the start timestamp, only running the method
     * if it is a different timestamp.
     */
    public static class CommitStageCondition<T> implements PreCommitCondition {
        private final BiFunction<Long, Long, T> timestampFunction;
        private final AtomicReference<T> commitStageResult;
        private volatile Optional<Long> startTimestamp;

        public CommitStageCondition(BiFunction<Long, Long, T> timestampFunction) {
            this.timestampFunction = timestampFunction;
            this.commitStageResult = new AtomicReference<>();
            this.startTimestamp = Optional.empty();
        }

        public void initialiseWithStartTimestamp(long startTs) {
            this.startTimestamp = Optional.of(startTs);
        }

        public T getCommitStageResult() {
            return commitStageResult.get();
        }

        @Override
        public void throwIfConditionInvalid(long timestamp) {
            Preconditions.checkState(startTimestamp.isPresent(), "Must initialise start timestamp immediately");
            long startTs = startTimestamp.get();

            if (startTs != timestamp) {
                this.commitStageResult.set(timestampFunction.apply(startTs, timestamp));
            }
        }
    }
}
